'''
This is a visual script that helps to position head frames into a sprite.

@author: kiira
'''
import pygame, os
from copy import copy

''' Change alpha for surfaces with per-pixel alpha; only for small surfaces '''
def change_alpha( surface, alpha=0.5 ):
    surface = copy( surface )
    size = surface.get_size()
    try:
        for y in xrange( size[1] ):
            for x in xrange( size[0] ):
                r, g, b, a = surface.get_at( ( x, y ) )
                surface.set_at( ( x, y ), ( r, g, b, int( a * alpha ) ) )
    except:
        return surface
    return surface

class ImageList:   
    
    def __init__( self, images, size, pos=( 0, 0 ), initial_index=0, resize=False ):
        self.size = size
        self.pos = pos
        
        self.arrow_size = ( size[0], int( size[1] * 0.05 ) )
        self.arrow_up_pos = self.pos
        self.arrow_down_pos = self.pos
        
        self.arrow_up = pygame.transform.smoothscale( pygame.image.load( 'arrow_up.png' ), self.arrow_size ).convert_alpha()
        self.translucent_arrow_up = change_alpha( self.arrow_up )
        self.arrow_down = pygame.transform.smoothscale( pygame.image.load( 'arrow_down.png' ), self.arrow_size ).convert_alpha()
        self.translucent_arrow_down = change_alpha( self.arrow_down )
        
          
        self.first_index = initial_index
        self.last_index = 0
        
        #reescale the images if necesary and add them to the images list
        self.images = []
        
        if resize:
            for name, image in images.iteritems():
                if image.get_width() != size[0]:
                    scale = size[0] / float( image.get_width() )
                    image = pygame.transform.smoothscale( image, ( int ( image.get_width() * scale ), int( image.get_height() * scale ) ) )
                    
                    self.images.append( [name, image, None] )
        else:
            for name, image in images.iteritems():
                self.images.append( [name, image, None] )       
            
        #update the showed frames
        self.update()
    
    def update( self ):       
        #aux variables
        broken = False
        xpos = self.pos[0]
        ypos = self.pos[1] + self.arrow_size[1] + 5
        
        #define the position only for the 
        for index in xrange( self.first_index, len( self.images ) ):
            current = self.images[index][1]
            
            if ypos + current.get_height() < self.pos[1] + self.size[1] - self.arrow_size[1]:
                self.images[index][2] = ( xpos, ypos )
                ypos += current.get_height()
            else:
                self.last_index = index
                broken = True
                break;
        
        if not broken:
            self.last_index = len( self.images )
                        
        #update arrow_down_pos
        self.arrow_down_pos = ( self.pos[0], self.pos[1] + self.size[1] - self.arrow_size[1] )           
        
    def render( self, surface ):           
        if self.first_index == 0:      
            surface.blit( self.translucent_arrow_up, self.arrow_up_pos )
        else:
            surface.blit( self.arrow_up, self.arrow_up_pos )
        
        for index in xrange( self.first_index, self.last_index ):   
            surface.blit( self.images[index][1], self.images[index][2] )
             
        if self.last_index == len( self.images ):
            surface.blit( self.translucent_arrow_down, self.arrow_down_pos, )
        else:
            surface.blit( self.arrow_down, self.arrow_down_pos )
    
    def react( self, point ):
        if point[0] > self.arrow_up_pos[0] and point[0] < self.arrow_up_pos[0] + self.arrow_size[0]:
            if point[1] > self.arrow_up_pos[1] and point[1] < self.arrow_up_pos[1] + self.arrow_size[1]:
                if self.first_index > 0:
                    self.first_index = self.first_index - 1      
                    self.update()
        
        if point[0] > self.arrow_down_pos[0] and point[0] < self.arrow_down_pos[0] + self.arrow_size[0]:
            if point[1] > self.arrow_down_pos[1] and point[1] < self.arrow_down_pos[1] + self.arrow_size[1]:
                if self.last_index < len( self.images ):
                    self.first_index = self.first_index + 1
                    self.update()
                    
        if point[0] > self.arrow_up_pos[0] and point[0] < self.arrow_up_pos[0] + self.arrow_size[0]:
            if point[1] > self.arrow_up_pos[1] + self.arrow_size[1] + 5 and point[1] < self.arrow_down_pos[1]:
                return self._get_image_at( point )           
            
    def _get_image_at( self, point ):
        for index in xrange( self.first_index, self.last_index ):
            if point[1] > self.images[index][2][1] and point[1] < self.images[index][2][1] + self.images[index][1].get_width():
                return copy( self.images[index] )  
            
class Button:
    def __init__( self, image, size, pos=( 0, 0 ), enabled=True ):
        if isinstance( image, pygame.Surface ):
            self.image = pygame.transform.smoothscale( image, size )
        else:
            self.image = pygame.transform.smoothscale( pygame.image.load( image ), size )
            
        self.translucent = change_alpha( self.image )
        self.pos = pos
        
        self.enabled = enabled
        
    def clicked( self, coords ):
        if self.enabled and self.image.get_rect( topleft=self.pos ).collidepoint( coords ):
            return True
    
    def enable( self ):
        self.enabled = True
        
    def disable( self ):
        self.enabled = False
    
    def render( self, surface ):
        if self.enabled:
            surface.blit( self.image, self.pos )
        else:
            surface.blit( self.translucent, self.pos )           
           
if __name__ == '__main__':
    from pygame.locals import QUIT, MOUSEBUTTONDOWN, MOUSEBUTTONUP, MOUSEMOTION, KEYDOWN, K_DOWN, K_UP, K_LEFT, K_RIGHT
    from Graphics import GraphicLoader
    from xml.dom import minidom
    from xml.dom.ext import PrettyPrint
    import operator
    
    
    #initialize pygame
    pygame.init()
    pygame.key.set_repeat( 50, 30 )
    
    #create the xml document
    if os.path.exists( 'head_positions.xml' ):
        head_pos_doc = minidom.parse( 'head_positions.xml' )
        positions = head_pos_doc.getElementsByTagName( 'positions' )[0]
    else:
        head_pos_doc = minidom.Document()
        positions = head_pos_doc.createElement( 'positions' )
        head_pos_doc.appendChild( positions )
            
    #initialization input
    while True:
        head_directory = raw_input( "Select the head directory to use: " ) 
        
        if os.path.isdir( head_directory ):
            break
        else:
            print 'Invalid directory.'
        
    while True:    
        sprite_directory = raw_input( "Select the sprite directory to use: " )
        
        if os.path.isdir( sprite_directory ):
            break
        else:
            print 'Invalid directory.'
            
    while True:
        try:
            scale = float( raw_input ( 'Select the scale of the sprites: ' ) )
            
            break;
        except:
            pass
        
    #initialize screen
    pantalla = pygame.display.set_mode( ( 400, 290 ) )    
    
    #create the loader for the sprite directory  
    sprite_loader = GraphicLoader( sprite_directory, "sprite" )
    current_sprite = None

    #create and populate the image dictionary
    file_names = os.listdir( head_directory )
    file_names.remove( '.svn' )
    image_dic = {}
        
    for file_name in file_names:   
        name = file_name.replace( '.png', '' )     
        image = pygame.image.load( os.path.join( head_directory, file_name ) )
        image = pygame.transform.smoothscale( image, ( int( image.get_width() * scale ), int( image.get_height() * scale ) ) )
        
        image_dic[name] = image
        
        if 'left' in name:
            reflected = name.replace( 'left', 'right' ) 
            
            image = pygame.transform.flip( pygame.image.load( os.path.join( head_directory, file_name ) ), 1, 0 )
            image = pygame.transform.smoothscale( image, ( int( image.get_width() * scale ), int( image.get_height() * scale ) ) )
            
            image_dic[reflected] = image
        if 'right' in name:
            reflected = name.replace( 'right', 'left' ) 
            
            image = pygame.transform.flip( pygame.image.load( os.path.join( head_directory, file_name ) ), 1, 0 )
            image = pygame.transform.smoothscale( image, ( int( image.get_width() * scale ), int( image.get_height() * scale ) ) )
            
            image_dic[reflected] = image
               
    #create the imagelist
    image_list = ImageList( image_dic, ( 30, 285 ) )
    
    #load the buttons
    ok_button = Button( "ok_button.png", ( 40, 40 ), ( 150, 250 ), False )
    open_button = Button( "open_button.png", ( 40, 40 ), ( 200, 250 ) )
    cancel_button = Button( "cancel_button.png", ( 40, 40 ), ( 250, 250 ), False )
    

    #aux variables
    selected = None
    drag = False
    finished = False
    front = True
                    
    while True:
        #rendering fase
        pantalla.fill( ( 255, 255, 255 ) )               
        image_list.render( pantalla )
        ok_button.render( pantalla )
        cancel_button.render( pantalla )                     
        open_button.render( pantalla )
        
        if front:
            if current_sprite:
                current_sprite.render( pantalla, ( 180, 80 ) ) 
            if selected:
                pantalla.blit( selected[1], selected[2] ) 
        else:   
            if selected:
                pantalla.blit( selected[1], selected[2] )       
            if current_sprite:
                current_sprite.render( pantalla, ( 180, 80 ) )             
                
        pygame.display.flip()
        
        #event handling
        for event in pygame.event.get():
            if event.type == QUIT and not current_sprite:
                #print the xml to a file
                archivo = open( "head_positions.xml", 'w' )
                PrettyPrint( head_pos_doc, stream=archivo, encoding="UTF-8" )
                archivo.close()
                
                #quit pygame and python
                pygame.quit()
                exit()
                
            elif event.type == MOUSEBUTTONDOWN and event.button == 1 :
                mouse_coords = pygame.mouse.get_pos()
                
                #image list event handler
                reaction = image_list.react( mouse_coords )
                    
                if reaction:
                    selected = reaction
                    drag = True
                elif selected and selected[1].get_rect( topleft=selected[2] ).collidepoint( mouse_coords ):
                    selected[2] = tuple( map( operator.sub, mouse_coords, tuple( map( operator.div, selected[1].get_size(), ( 2 , 2 ) ) ) ) )
                    drag = True               
                    
                #buttons event handler
                elif open_button.clicked( mouse_coords ):                                                                                
                    
                    try:
                        #load the sprite
                        action = raw_input( "Select the action sprite to load: " )
                        anchor = raw_input( "Anchor: " )                    
                        
                        current_sprite = sprite_loader.load( action, anchor=anchor, scale=( scale, scale ) )
                                               
                        #aux variables for node creation
                        looking_at = raw_input( "Looking at? " )
                        frames = {}
                        finished = False
                        
                        #button states
                        open_button.disable()
                        ok_button.enable()
                        cancel_button.enable()
                        
                    except RuntimeError:
                        print "Sprite not found."
                                                                        
                elif ok_button.clicked( mouse_coords ):
                    if not finished and selected:
                        #set the current frame image and position
                        frames[current_sprite.get_current_frame()] = ( selected[1], selected[0], map( operator.sub, selected[2], ( 180, 80 ) ), front ) 
                    
                        current_sprite.next_frame()
                    
                        if current_sprite.get_current_frame() == 0:
                            selected = None
                        
                            #load the attachs and add them to the sprite
                            for i, frame in frames.iteritems():
                                current_sprite.attach( frame[0], i, frame[2], frame[3] )
                            
                            current_sprite.play()  
                            
                            finished = True
                    
                    elif finished:
                        #removes the action node for the current sprite if it exists
                        actions = positions.getElementsByTagName( 'action' )
                        
                        for node in actions:
                            if node.getAttribute( 'name' ) == action and node.getAttribute( 'looking_at' ) == looking_at:
                                positions.removeChild( node )
                                
                        #creates and modifies the action node for the current action
                        action_node = head_pos_doc.createElement( 'action' )
                        action_node.setAttribute( 'looking_at', looking_at )
                        action_node.setAttribute( 'name', action )

                        
                        #create and append frame nodes
                        for key, frame in frames.iteritems():  
                            #modify the image pos to reflect the center of the image instead of the leftop corner
                            half_x = tuple( map( operator.div, frame[0].get_size(), ( 2, 1 ) ) )
                            pos = tuple( map( operator.add, frame[2], half_x ) )
                            
                            
                            frame_node = head_pos_doc.createElement( 'frame' )
                            
                            if not frame[3]:
                                frame_node.setAttribute( 'front', 'False' )
                            
                            frame_node.setAttribute( 'number', str( key ) )
                                                                                                                
                            image_node = head_pos_doc.createElement( 'image' )
                            text_node = head_pos_doc.createTextNode( frame[1] )
                            image_node.appendChild( text_node )
                            frame_node.appendChild( image_node )
                            
                            xpos_node = head_pos_doc.createElement( 'xpos' )
                            text_node = head_pos_doc.createTextNode( str( pos[0] ) )
                            xpos_node.appendChild( text_node )
                            frame_node.appendChild( xpos_node )
                            
                            ypos_node = head_pos_doc.createElement( 'ypos' )
                            text_node = head_pos_doc.createTextNode( str( pos[1] ) )
                            ypos_node.appendChild( text_node )
                            frame_node.appendChild( ypos_node )
                            
                            action_node.appendChild( frame_node )
                            
                        positions.appendChild( action_node )
                        
                        current_sprite = None
                        finished = False                        
                        
                        open_button.enable()
                        ok_button.disable()
                        cancel_button.disable()  
                
                elif cancel_button.clicked( mouse_coords ):
                    current_sprite = None  
                      
                    open_button.enable()
                    ok_button.disable()
                    cancel_button.disable()               
            
            elif event.type == MOUSEBUTTONDOWN and event.button == 3:
                front = not front
                                                                                     
            elif event.type == MOUSEBUTTONUP and event.button == 1:
                drag = False            
            elif event.type == MOUSEMOTION and drag:
                selected[2] = tuple( map( operator.sub, pygame.mouse.get_pos(), tuple( map( operator.div, selected[1].get_size(), ( 2 , 2 ) ) ) ) )
            elif event.type == KEYDOWN and selected:
                if event.key == K_DOWN:
                    selected[2] = tuple( map( operator.add, selected[2], ( 0, 1 ) ) )
                elif event.key == K_UP:
                    selected[2] = tuple( map( operator.add, selected[2], ( 0, -1 ) ) )
                elif event.key == K_RIGHT:
                    selected[2] = tuple( map( operator.add, selected[2], ( 1, 0 ) ) )
                elif event.key == K_LEFT:
                    selected[2] = tuple( map( operator.add, selected[2], ( -1, 0 ) ) )
